import rospy
import os
import cv2
import json
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import threading
import tf2_ros
import geometry_msgs.msg

class DataCollector:
    def __init__(self):
        # 存储根目录配置
        self.data_root = os.path.expanduser("~/Documents/collecting_data")
        self.collection_number = self.find_latest_collection_number() + 1
        
        # 同步控制
        self.recording = False
        self.image_lock = threading.Lock()
        self.bridge = CvBridge()
        
        # 相机数据缓存
        self.camera_data = {
            "blot": {"rgb": None, "depth": None, "ir": None},
            "realsense": {"rgb": None, "depth": None}
        }
        
        # 初始化订阅
        self.init_subscribers()
        
        # TF监听器
        self.tf_buffer = tf2_ros.Buffer(cache_time=rospy.Duration(10))
        self.listener = tf2_ros.TransformListener(self.tf_buffer)
        
        print("等待指令：'s'-开始 'e'-结束 'q'-退出")

    def find_latest_collection_number(self):
        if not os.path.exists(self.data_root):
            return 0
        folders = [f for f in os.listdir(self.data_root) if f.isdigit()]
        return max([int(f) for f in folders]) if folders else 0

    def init_subscribers(self):
        # Blot相机
        rospy.Subscriber("/femto/color/image_raw", Image, 
                       self.image_callback, callback_args=("blot", "rgb"))
        rospy.Subscriber("/femto/depth/image_raw", Image,
                       self.image_callback, callback_args=("blot", "depth"))
        rospy.Subscriber("/femto/ir/image_raw", Image,
                       self.image_callback, callback_args=("blot", "ir"))
        
        # Realsense
        rospy.Subscriber("/realsense/color/image_raw", Image,
                       self.image_callback, callback_args=("realsense", "rgb"))
        rospy.Subscriber("/realsense/depth/image_rect_raw", Image,
                       self.image_callback, callback_args=("realsense", "depth"))

    def image_callback(self, msg, camera_sensor):
        camera, sensor = camera_sensor
        if self.recording:
            with self.image_lock:
                try:
                    # 转换图像格式
                    if sensor == "rgb":
                        cv_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")
                    else:
                        cv_image = self.bridge.imgmsg_to_cv2(msg, "8UC1")
                    
                    # 更新缓存（带时间戳）
                    self.camera_data[camera][sensor] = (cv_image, msg.header.stamp)
                    
                except Exception as e:
                    rospy.logerr(f"图像转换失败 ({camera}/{sensor}): {str(e)}")

    def create_storage_folders(self):
        """创建并列的三个存储目录"""
        base_path = os.path.join(self.data_root, f"{self.collection_number:04d}")
        
        # 主要目录结构
        self.blot_path = os.path.join(base_path, "blot")
        self.realsense_path = os.path.join(base_path, "realsense")
        self.pose_path = os.path.join(base_path, "poses")
        
        # 创建目录
        os.makedirs(self.blot_path, exist_ok=True)
        os.makedirs(self.realsense_path, exist_ok=True)
        os.makedirs(self.pose_path, exist_ok=True)

    def save_camera_images(self, camera_type):
        """保存指定相机的图像数据"""
        save_path = self.blot_path if camera_type == "blot" else self.realsense_path
        data = self.camera_data[camera_type]
        saved_stamps = []
        
        for sensor in data:
            if data[sensor] is not None:
                img, stamp = data[sensor]
                timestamp_str = f"{stamp.to_sec():.6f}".replace('.', '_')
                
                # 保存图像
                filename = os.path.join(save_path, f"{sensor}_{timestamp_str}.png")
                cv2.imwrite(filename, img)
                rospy.logdebug(f"保存 {camera_type}/{sensor}: {filename}")
                saved_stamps.append(stamp)
                data[sensor] = None  # 清空缓存
        
        return max(saved_stamps) if saved_stamps else None

    def save_pose_data(self, timestamp):
        """保存位姿数据到poses目录"""
        try:
            transform = self.tf_buffer.lookup_transform(
                'panda_link0', 
                'panda_link8',
                timestamp,
                timeout=rospy.Duration(0.1))
            
            pose_data = {
                "timestamp": timestamp.to_sec(),
                "position": {
                    "x": transform.transform.translation.x,
                    "y": transform.transform.translation.y,
                    "z": transform.transform.translation.z
                },
                "orientation": {
                    "x": transform.transform.rotation.x,
                    "y": transform.transform.rotation.y,
                    "z": transform.transform.rotation.z,
                    "w": transform.transform.rotation.w
                }
            }
            
            timestamp_str = f"{timestamp.to_sec():.6f}".replace('.', '_')
            filename = os.path.join(self.pose_path, f"pose_{timestamp_str}.json")
            with open(filename, 'w') as f:
                json.dump(pose_data, f)
            rospy.logdebug(f"保存位姿: {filename}")
            
        except (tf2_ros.LookupException,
                tf2_ros.ConnectivityException,
                tf2_ros.ExtrapolationException) as e:
            rospy.logwarn(f"位姿获取失败: {str(e)}")

    def record_data(self, event):
        """定时采集回调（5Hz）"""
        with self.image_lock:
            # 复制数据并清空缓存
            blot_data = self.camera_data["blot"].copy()
            realsense_data = self.camera_data["realsense"].copy()

            # 清空缓存
            # for sensor in self.camera_data["blot"]:
            #     self.camera_data["blot"][sensor] = None
            # for sensor in self.camera_data["realsense"]:
            #     self.camera_data["realsense"][sensor] = None
        
        # 保存数据并获取时间戳
        blot_stamp = self.save_camera_images("blot")
        realsense_stamp = self.save_camera_images("realsense")
        
        # 使用最新时间戳保存位姿
        latest_stamp = max(filter(None, [blot_stamp, realsense_stamp]), default=None)
        if latest_stamp:
            self.save_pose_data(latest_stamp)

    def start_recording(self):
        if not self.recording:
            self.recording = True
            self.create_storage_folders()
            self.timer = rospy.Timer(rospy.Duration(1), self.record_data)  # 5Hz
            rospy.loginfo(f"开始采集 #{self.collection_number}")

    def stop_recording(self):
        if self.recording:
            self.recording = False
            if self.timer:
                self.timer.shutdown()
            self.collection_number += 1
            rospy.loginfo(f"采集完成: {self.data_root}/{self.collection_number-1:04d}")

    def run(self):
        while not rospy.is_shutdown():
            cmd = input().lower()
            if cmd == 's':
                self.start_recording()
            elif cmd == 'e':
                self.stop_recording()
            elif cmd == 'q':
                rospy.signal_shutdown("用户退出")
                break

if __name__ == "__main__":
    rospy.init_node('multi_camera_collector', log_level=rospy.INFO)
    collector = DataCollector()
    collector.run()